{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ae744dd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from math import sqrt\n",
    "import copy\n",
    "import  traceback\n",
    "import shutil\n",
    "\n",
    "import numpy as np  # linear algebra\n",
    "import pydicom\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "from pydicom.uid import UID\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3d011f94",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_scan(path):\n",
    "    slices = [pydicom.dcmread(path + '/' + s) for s in filter(lambda x: x.endswith('.dcm'), os.listdir(path))]\n",
    "    slices.sort(key=lambda x: float(x.InstanceNumber))\n",
    "    return slices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "043bd268",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "yujinmin-Im25-78\n"
     ]
    }
   ],
   "source": [
    "cta_root_path = '/nfs3-p1/zsxm/dataset/aorta_img_label/'\n",
    "root_path = '/nfs3-p1/zsxm/dataset/aorta/'\n",
    "input_folder_list = []\n",
    "for path in os.listdir(root_path):\n",
    "    if not os.path.isdir(os.path.join(root_path, path)):\n",
    "        continue\n",
    "    if os.path.exists(os.path.join(root_path, path, '1/')) and os.path.exists(os.path.join(root_path, path, '2/')):\n",
    "        input_folder_list.append(os.path.join(root_path, path))\n",
    "    else:\n",
    "        print(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b222e8c5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing gengyongchao-Im21-83\n",
      "Processing hebin-Im41-91\n",
      "Ignore hebin-Im41-91\n",
      "Processing hejianjun-Im31-99\n",
      "Processing heshengnan-Im16-90\n",
      "Processing huangjianhua-Im38-144\n",
      "Processing hujiaguo-Im34-58\n",
      "Processing jiangquanfeng-Im 50-107\n",
      "Processing litingxian-Im36-141\n",
      "Processing luohuafa-im29-96\n",
      "Ignore luohuafa-im29-96\n",
      "Processing lushida-Im33-159\n",
      "Processing lvpeiji-Im32-136\n",
      "Processing manyanfeng-Im32-153\n",
      "Processing maoliang-Im12-90\n",
      "Processing modehe-Im34-89\n",
      "Processing nibingshan-Im33-42\n",
      "Processing qiquanping-Im25-54\n",
      "Processing shengfumei-Im31-132\n",
      "Processing shenhuagou-Im33-95\n",
      "Processing shenjiaren-Im35-136\n",
      "Processing shenjie-Im36-100\n",
      "Processing shiaihua-Im28-108\n",
      "Ignore shiaihua-Im28-108\n",
      "Processing shilei-Im21-92\n",
      "Ignore shilei-Im21-92\n",
      "Processing shuzhiming-Im28-80\n",
      "Processing taofangguo-Im30-142\n",
      "Processing wangaishui-Im75-128\n",
      "Processing wanggangxiang-Im16-90\n",
      "Processing wanggaosheng-Im24-57\n",
      "Ignore wanggaosheng-Im24-57\n",
      "Processing wangjianjun-Im30-110\n",
      "Processing wanglianlian-Im42-139\n",
      "Processing wangmanying-Im23-134\n",
      "Processing wangzhongwen-Im36-80\n",
      "Ignore wangzhongwen-Im36-80\n",
      "Processing wuguangde-im55-205\n",
      "Processing xiewenbo-Im25-40\n",
      "Processing xujinbiao-Im37-146\n",
      "Processing xujinyang-Im15-30\n",
      "Processing yangwenbin-Im44-125\n",
      "Processing yangyoufu-Im16-25\n",
      "Processing yaoliang-Im30-83\n",
      "Processing yinghao-Im18-86\n",
      "Ignore yinghao-Im18-86\n",
      "Processing yingzhemin-Im58-108\n",
      "Processing yuanenchun-Im27-93\n",
      "Processing yukongban-Im30-139\n",
      "Processing yuliyao-Im25-31\n",
      "Processing zengzhengyou-Im54-145\n",
      "Processing zhangxinliang-Im16-24\n",
      "Processing zhangxiyuan-Im9-50\n",
      "Processing zhaojianbin-Im31-158\n",
      "Processing zhaojun-Im45-103\n",
      "Processing zhengsiqing-Im35-106\n",
      "Processing zhoufei-Im16-97\n",
      "Processing zhoujinhai-Im14-94\n",
      "Processing chuguangzhou-Im46-111\n",
      "Ignore chuguangzhou-Im46-111\n",
      "Processing dongxuesheng-Im37-149\n",
      "Processing durenxi-Im57-172\n",
      "Processing fangcaiyu-Im21-91\n",
      "Processing fengyonghua-Im21-89\n",
      "Processing fuhuajin-Im29-126\n",
      "Processing gaomingsheng-Im34-151\n",
      "Processing zhoujun-Im33-128\n",
      "Processing zhouzheyu-Im32-152\n",
      "Processing zhudehong-Im22-91\n",
      "Processing zhuyongbu-im23-92\n",
      "Ignore zhuyongbu-im23-92\n",
      "Processing gengguorui-Im18-66\n",
      "Ignore gengguorui-Im18-66\n",
      "Processing chengxuejun-Im17-83\n",
      "Ignore chengxuejun-Im17-83\n"
     ]
    }
   ],
   "source": [
    "complete_path = []\n",
    "ignore_path = []\n",
    "for patient in input_folder_list:\n",
    "    sl = patient.split('/')\n",
    "    name = sl[-2] if sl[-1] == '' else sl[-1]\n",
    "    print(f'Processing {name}')\n",
    "    \n",
    "    image_path = os.path.join(patient, '1', 'images')\n",
    "    label_path = os.path.join(patient, '1', 'labels')\n",
    "    if os.path.exists(image_path):\n",
    "        shutil.rmtree(image_path)\n",
    "    if os.path.exists(label_path):\n",
    "        shutil.rmtree(label_path)\n",
    "    os.mkdir(image_path)\n",
    "    os.mkdir(label_path)\n",
    "    with open(os.path.join(label_path, f'labels.txt'), 'w') as txt:\n",
    "        txt.write('aorta')\n",
    "    \n",
    "    ct = load_scan(os.path.join(patient, '1'))\n",
    "    cta = load_scan(os.path.join(patient, '2'))\n",
    "    try:\n",
    "        ct_head, ct_tail, cta_head, cta_tail = ct[0].ImagePositionPatient[2],ct[-1].ImagePositionPatient[2],cta[0].ImagePositionPatient[2],cta[-1].ImagePositionPatient[2]\n",
    "    except:\n",
    "        ct_head, ct_tail, cta_head, cta_tail = ct[0].SliceLocation,ct[-1].SliceLocation,cta[0].SliceLocation,cta[-1].SliceLocation\n",
    "    if ct_head > cta_head or ct_tail < cta_tail:\n",
    "        print(f'Ignore {name}')\n",
    "        ignore_path.append(patient)\n",
    "        shutil.rmtree(image_path)\n",
    "        shutil.rmtree(label_path)\n",
    "        continue\n",
    "    \n",
    "    p = 0\n",
    "    for i in range(len(ct)):\n",
    "        img = ct[i].pixel_array.astype(np.int16)\n",
    "        intercept = ct[i].RescaleIntercept\n",
    "        slope = ct[i].RescaleSlope\n",
    "        if slope != 1:\n",
    "            img = (slope * img.astype(np.float64)).astype(np.int16)\n",
    "        img += np.int16(intercept)\n",
    "        img = np.clip(img, -100, 500)\n",
    "        img = ((img+100)/600*255).astype(np.uint8)\n",
    "        img = Image.fromarray(img)\n",
    "        img.save(os.path.join(image_path, f'{name}_{i:04d}.png'))\n",
    "        with open(os.path.join(label_path, f'{name}_{i:04d}.txt'), 'w') as txt:\n",
    "            while p+1 < len(cta):\n",
    "                try:\n",
    "                    ct_z, cta_p_z, cta_p1_z = ct[i].ImagePositionPatient[2], cta[p].ImagePositionPatient[2], cta[p+1].ImagePositionPatient[2]\n",
    "                except:\n",
    "                    ct_z, cta_p_z, cta_p1_z = ct[i].SliceLocation, cta[p].SliceLocation, cta[p+1].SliceLocation\n",
    "                if (cta_p_z-ct_z) * (cta_p1_z-ct_z) <= 0:\n",
    "                    p = p if abs(cta_p_z-ct_z) < abs(cta_p1_z-ct_z) else p+1\n",
    "                    with open(os.path.join(cta_root_path, name, 'labels', f'{name}_{p:04d}.txt'), 'r') as sf:\n",
    "                        txt.writelines(sf.readlines())\n",
    "                    break\n",
    "                p += 1\n",
    "            else:\n",
    "                raise Exception(f'CT图像范围超出了CTA图像范围：{ct_z},{cta_p_z},{cta_p1_z}')\n",
    "    complete_path.append(patient)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "27193ddf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "64 53 11\n"
     ]
    }
   ],
   "source": [
    "print(len(input_folder_list), len(complete_path), len(ignore_path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4f068b6d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing hebin-Im41-91\n",
      "Processing luohuafa-im29-96\n",
      "Processing shiaihua-Im28-108\n",
      "Processing shilei-Im21-92\n",
      "Processing wanggaosheng-Im24-57\n",
      "Processing wangzhongwen-Im36-80\n",
      "Processing yinghao-Im18-86\n",
      "Processing chuguangzhou-Im46-111\n",
      "Processing zhuyongbu-im23-92\n",
      "Processing gengguorui-Im18-66\n",
      "Processing chengxuejun-Im17-83\n"
     ]
    }
   ],
   "source": [
    "for patient in ignore_path:\n",
    "    sl = patient.split('/')\n",
    "    name = sl[-2] if sl[-1] == '' else sl[-1]\n",
    "    print(f'Processing {name}')\n",
    "    \n",
    "    image_path = os.path.join(patient, '1', 'images')\n",
    "    if os.path.exists(image_path):\n",
    "        shutil.rmtree(image_path)\n",
    "    os.mkdir(image_path)\n",
    "    \n",
    "    ct = load_scan(os.path.join(patient, '1'))\n",
    "    \n",
    "    for i in range(len(ct)):\n",
    "        img = ct[i].pixel_array.astype(np.int16)\n",
    "        intercept = ct[i].RescaleIntercept\n",
    "        slope = ct[i].RescaleSlope\n",
    "        if slope != 1:\n",
    "            img = (slope * img.astype(np.float64)).astype(np.int16)\n",
    "        img += np.int16(intercept)\n",
    "        img = np.clip(img, -100, 500)\n",
    "        img = ((img+100)/600*255).astype(np.uint8)\n",
    "        img = Image.fromarray(img)\n",
    "        img.save(os.path.join(image_path, f'{name}_{i:04d}.png'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "84ce50e2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Copy /nfs3-p1/zsxm/dataset/aorta/hebin-Im41-91 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/wrong/hebin-Im41-91\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/luohuafa-im29-96 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/wrong/luohuafa-im29-96\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/shiaihua-Im28-108 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/wrong/shiaihua-Im28-108\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/shilei-Im21-92 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/wrong/shilei-Im21-92\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/wanggaosheng-Im24-57 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/wrong/wanggaosheng-Im24-57\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/wangzhongwen-Im36-80 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/wrong/wangzhongwen-Im36-80\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/yinghao-Im18-86 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/wrong/yinghao-Im18-86\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/chuguangzhou-Im46-111 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/wrong/chuguangzhou-Im46-111\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/zhuyongbu-im23-92 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/wrong/zhuyongbu-im23-92\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/gengguorui-Im18-66 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/wrong/gengguorui-Im18-66\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/chengxuejun-Im17-83 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/wrong/chengxuejun-Im17-83\n"
     ]
    }
   ],
   "source": [
    "dst_path = '/nfs3-p1/zsxm/dataset/aorta_ct_img_label/wrong'\n",
    "if not os.path.exists(dst_path):\n",
    "    os.mkdir(dst_path)\n",
    "for p in ignore_path:\n",
    "    sl = p.split('/')\n",
    "    name = sl[-1] if sl[-1] != '' else sl[-2]\n",
    "    print(f\"Copy {os.path.join(p)} To {os.path.join(dst_path, name)}\")\n",
    "    if os.path.exists(os.path.join(dst_path, name)):\n",
    "        print(f\"remove {os.path.join(dst_path, name)}\")\n",
    "        shutil.rmtree(os.path.join(dst_path, name))\n",
    "    os.mkdir(os.path.join(dst_path, name))\n",
    "        \n",
    "    try:\n",
    "        shutil.copytree(os.path.join(p, '1', 'images'), os.path.join(dst_path, name, 'images'))\n",
    "    except:\n",
    "        traceback.print_exc()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "498a2420",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Copy /nfs3-p1/zsxm/dataset/aorta/gengyongchao-Im21-83 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/gengyongchao-Im21-83\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/hejianjun-Im31-99 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/hejianjun-Im31-99\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/heshengnan-Im16-90 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/heshengnan-Im16-90\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/huangjianhua-Im38-144 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/huangjianhua-Im38-144\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/hujiaguo-Im34-58 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/hujiaguo-Im34-58\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/jiangquanfeng-Im 50-107 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/jiangquanfeng-Im 50-107\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/litingxian-Im36-141 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/litingxian-Im36-141\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/lushida-Im33-159 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/lushida-Im33-159\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/lvpeiji-Im32-136 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/lvpeiji-Im32-136\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/manyanfeng-Im32-153 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/manyanfeng-Im32-153\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/maoliang-Im12-90 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/maoliang-Im12-90\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/modehe-Im34-89 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/modehe-Im34-89\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/nibingshan-Im33-42 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/nibingshan-Im33-42\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/qiquanping-Im25-54 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/qiquanping-Im25-54\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/shengfumei-Im31-132 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/shengfumei-Im31-132\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/shenhuagou-Im33-95 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/shenhuagou-Im33-95\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/shenjiaren-Im35-136 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/shenjiaren-Im35-136\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/shenjie-Im36-100 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/shenjie-Im36-100\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/shuzhiming-Im28-80 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/shuzhiming-Im28-80\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/taofangguo-Im30-142 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/taofangguo-Im30-142\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/wangaishui-Im75-128 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/wangaishui-Im75-128\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/wanggangxiang-Im16-90 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/wanggangxiang-Im16-90\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/wangjianjun-Im30-110 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/wangjianjun-Im30-110\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/wanglianlian-Im42-139 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/wanglianlian-Im42-139\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/wangmanying-Im23-134 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/wangmanying-Im23-134\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/wuguangde-im55-205 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/wuguangde-im55-205\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/xiewenbo-Im25-40 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/xiewenbo-Im25-40\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/xujinbiao-Im37-146 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/xujinbiao-Im37-146\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/xujinyang-Im15-30 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/xujinyang-Im15-30\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/yangwenbin-Im44-125 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/yangwenbin-Im44-125\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/yangyoufu-Im16-25 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/yangyoufu-Im16-25\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/yaoliang-Im30-83 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/yaoliang-Im30-83\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/yingzhemin-Im58-108 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/yingzhemin-Im58-108\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/yuanenchun-Im27-93 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/yuanenchun-Im27-93\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/yukongban-Im30-139 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/yukongban-Im30-139\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/yuliyao-Im25-31 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/yuliyao-Im25-31\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/zengzhengyou-Im54-145 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/zengzhengyou-Im54-145\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/zhangxinliang-Im16-24 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/zhangxinliang-Im16-24\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/zhangxiyuan-Im9-50 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/zhangxiyuan-Im9-50\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/zhaojianbin-Im31-158 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/zhaojianbin-Im31-158\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/zhaojun-Im45-103 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/zhaojun-Im45-103\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/zhengsiqing-Im35-106 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/zhengsiqing-Im35-106\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/zhoufei-Im16-97 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/zhoufei-Im16-97\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/zhoujinhai-Im14-94 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/zhoujinhai-Im14-94\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/dongxuesheng-Im37-149 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/dongxuesheng-Im37-149\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/durenxi-Im57-172 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/durenxi-Im57-172\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/fangcaiyu-Im21-91 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/fangcaiyu-Im21-91\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/fengyonghua-Im21-89 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/fengyonghua-Im21-89\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/fuhuajin-Im29-126 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/fuhuajin-Im29-126\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/gaomingsheng-Im34-151 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/gaomingsheng-Im34-151\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/zhoujun-Im33-128 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/zhoujun-Im33-128\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/zhouzheyu-Im32-152 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/zhouzheyu-Im32-152\n",
      "Copy /nfs3-p1/zsxm/dataset/aorta/zhudehong-Im22-91 To /nfs3-p1/zsxm/dataset/aorta_ct_img_label/new/zhudehong-Im22-91\n"
     ]
    }
   ],
   "source": [
    "dst_path = '/nfs3-p1/zsxm/dataset/aorta_ct_img_label/new'\n",
    "if not os.path.exists(dst_path):\n",
    "    os.mkdir(dst_path)\n",
    "for p in complete_path:\n",
    "    sl = p.split('/')\n",
    "    name = sl[-1] if sl[-1] != '' else sl[-2]\n",
    "    print(f\"Copy {os.path.join(p)} To {os.path.join(dst_path, name)}\")\n",
    "    if os.path.exists(os.path.join(dst_path, name)):\n",
    "        print(f\"remove {os.path.join(dst_path, name)}\")\n",
    "        shutil.rmtree(os.path.join(dst_path, name))\n",
    "    os.mkdir(os.path.join(dst_path, name))\n",
    "        \n",
    "    try:\n",
    "        shutil.copytree(os.path.join(p, '1', 'images'), os.path.join(dst_path, name, 'images'))\n",
    "        shutil.copytree(os.path.join(p, '1', 'labels'), os.path.join(dst_path, name, 'labels'))\n",
    "    except:\n",
    "        traceback.print_exc()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e4ef9998",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing dongxuesheng-Im37-149\n",
      "Processing durenxi-Im57-172\n",
      "Processing fangcaiyu-Im21-91\n",
      "Processing fengyonghua-Im21-89\n",
      "Processing fuhuajin-Im29-126\n",
      "Processing gaomingsheng-Im34-151\n",
      "Processing gengyongchao-Im21-83\n",
      "Processing hejianjun-Im31-99\n",
      "Processing heshengnan-Im16-90\n",
      "Processing huangjianhua-Im38-144\n",
      "Processing jiangquanfeng-Im 50-107\n",
      "Processing litingxian-Im36-141\n",
      "Processing lushida-Im33-159\n",
      "Processing manyanfeng-Im32-153\n",
      "Processing modehe-Im34-89\n",
      "Processing nibingshan-Im33-42\n",
      "Processing shenhuagou-Im33-95\n",
      "Processing shenjie-Im36-100\n",
      "Processing shuzhiming-Im28-80\n",
      "Processing taofangguo-Im30-142\n",
      "Processing wangaishui-Im75-128\n",
      "Processing wanggangxiang-Im16-90\n",
      "Processing wangjianjun-Im30-110\n",
      "Processing wanglianlian-Im42-139\n",
      "Processing wangmanying-Im23-134\n",
      "Processing yangwenbin-Im44-125\n",
      "Processing yaoliang-Im30-83\n",
      "Processing yuanenchun-Im27-93\n",
      "Processing yuliyao-Im25-31\n",
      "Processing zengzhengyou-Im54-145\n",
      "Processing zhangxinliang-Im16-24\n",
      "Processing zhaojianbin-Im31-158\n",
      "Processing zhoufei-Im16-97\n",
      "Processing zhudehong-Im22-91\n",
      "Processing wuguangde-im55-205\n",
      "Processing xiewenbo-Im25-40\n",
      "Processing xujinyang-Im15-30\n",
      "Processing yangyoufu-Im16-25\n",
      "Processing yingzhemin-Im58-108\n",
      "Processing yukongban-Im30-139\n",
      "Processing zhangxiyuan-Im9-50\n",
      "Processing zhaojun-Im45-103\n",
      "Processing zhoujinhai-Im14-94\n",
      "Processing zhoujun-Im33-128\n",
      "Processing zhouzheyu-Im32-152\n",
      "Processing lvpeiji-Im32-136\n",
      "Processing maoliang-Im12-90\n",
      "Processing shenjiaren-Im35-136\n",
      "Processing chengxuejun-Im17-83\n",
      "Processing chuguangzhou-Im46-111\n",
      "Processing gengguorui-Im18-66\n",
      "Processing hebin-Im41-91\n",
      "Processing luohuafa-im29-96\n",
      "Processing shiaihua-Im28-108\n",
      "Processing shilei-Im21-92\n",
      "Processing wanggaosheng-Im24-57\n",
      "Processing wangzhongwen-Im36-80\n",
      "Processing yinghao-Im18-86\n",
      "Processing zhuyongbu-im23-92\n",
      "Processing hujiaguo-Im34-58\n",
      "Processing shengfumei-Im31-132\n",
      "Processing xujinbiao-Im37-146\n",
      "Processing zhengsiqing-Im35-106\n",
      "Processing qiquanping-Im25-54\n"
     ]
    }
   ],
   "source": [
    "def move_to_yolo(input_path, img_path, label_path):\n",
    "    if not os.path.exists(img_path):\n",
    "        os.makedirs(img_path)\n",
    "    if not os.path.exists(label_path):\n",
    "        os.makedirs(label_path)\n",
    "\n",
    "    for patient in os.listdir(input_path):\n",
    "        if not os.path.isdir(os.path.join(input_path, patient)):\n",
    "            continue\n",
    "\n",
    "        print(f'Processing {patient}')\n",
    "\n",
    "        for f in os.listdir(os.path.join(input_path, patient, 'images')):\n",
    "            shutil.copy(os.path.join(input_path, patient, 'images',f), img_path)\n",
    "\n",
    "        for f in os.listdir(os.path.join(input_path, patient, 'labels')):\n",
    "            if f == 'labels.txt':\n",
    "                continue\n",
    "            shutil.copy(os.path.join(input_path, patient, 'labels',f), label_path)\n",
    "\n",
    "input_path = '/nfs3-p1/zsxm/dataset/aorta_ct_img_label/train/'\n",
    "img_path = '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/images/train/'\n",
    "label_path = '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/labels_ori/train/'\n",
    "move_to_yolo(input_path, img_path, label_path)\n",
    "input_path = '/nfs3-p1/zsxm/dataset/aorta_ct_img_label/val/'\n",
    "img_path = '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/images/val/'\n",
    "label_path = '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/labels_ori/val/'\n",
    "move_to_yolo(input_path, img_path, label_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "fe193cfb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def correct_to_square(input_path, output_path):\n",
    "    if not os.path.exists(output_path):\n",
    "        os.makedirs(output_path)\n",
    "\n",
    "    for txt in os.listdir(input_path):\n",
    "        if not txt.endswith('.txt') or txt == 'labels.txt':\n",
    "            continue\n",
    "\n",
    "        with open(os.path.join(input_path, txt), 'r') as sf:\n",
    "            lines = sf.readlines()\n",
    "        assert len(lines) < 3\n",
    "        with open(os.path.join(output_path, txt), 'w') as df:\n",
    "            for i in range(len(lines)):\n",
    "                aorta = lines[i].split()\n",
    "                aorta_f = list(map(lambda x: float(x), aorta))\n",
    "                max_edge = max(float(aorta[3]), float(aorta[4]))\n",
    "                head = \"\" if i == 0 else \"\\n\"\n",
    "                df.write(f'{head}0 {aorta[1]} {aorta[2]} {max_edge:.6f} {max_edge:.6f}')\n",
    "                \n",
    "input_path = '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/labels_ori/train/'\n",
    "output_path = '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/labels/train/'\n",
    "correct_to_square(input_path, output_path)\n",
    "input_path = '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/labels_ori/val/'\n",
    "output_path = '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/labels/val/'\n",
    "correct_to_square(input_path, output_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "19367402",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[], [], [], [], [], [], [], [], [], []]\n"
     ]
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cddb8cb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#这种做法是错误的，会导致candidate_table[0] == candidate_table[1] 改变其中一个另一个也跟着改变#candidate_table = [[]] * 256\n",
    "train_ct_path = '/nfs3-p1/zsxm/dataset/aorta_ct_img_label/train/'\n",
    "train_ct_list = os.listdir(train_ct_path)\n",
    "for patient in input_folder_list:\n",
    "    sl = patient.split('/')\n",
    "    name = sl[-2] if sl[-1] == '' else sl[-1]\n",
    "    if name not in train_ct_list:\n",
    "        continue\n",
    "    print(f'Processing {name}')\n",
    "    \n",
    "    ct = load_scan(os.path.join(patient, '1'))\n",
    "    cta = load_scan(os.path.join(patient, '2'))\n",
    "    try:\n",
    "        ct_head, ct_tail, cta_head, cta_tail = ct[0].ImagePositionPatient[2],ct[-1].ImagePositionPatient[2],cta[0].ImagePositionPatient[2],cta[-1].ImagePositionPatient[2]\n",
    "    except:\n",
    "        ct_head, ct_tail, cta_head, cta_tail = ct[0].SliceLocation,ct[-1].SliceLocation,cta[0].SliceLocation,cta[-1].SliceLocation\n",
    "    if ct_head > cta_head or ct_tail < cta_tail:\n",
    "        print(f'Ignore {name}')\n",
    "        continue\n",
    "    \n",
    "    p = 0\n",
    "    for i in range(len(ct)):\n",
    "        img = ct[i].pixel_array.astype(np.int16)\n",
    "        intercept = ct[i].RescaleIntercept\n",
    "        slope = ct[i].RescaleSlope\n",
    "        if slope != 1:\n",
    "            img = (slope * img.astype(np.float64)).astype(np.int16)\n",
    "        img += np.int16(intercept)\n",
    "        img = np.clip(img, -100, 500)\n",
    "        img = ((img+100)/600*255).astype(np.uint8)\n",
    "#         img = Image.fromarray(img)\n",
    "#         img.save(os.path.join(patient, '1', 'images', f'{name}_{i:04d}.png'))\n",
    "        while p+1 < len(cta):\n",
    "            try:\n",
    "                ct_z, cta_p_z, cta_p1_z = ct[i].ImagePositionPatient[2], cta[p].ImagePositionPatient[2], cta[p+1].ImagePositionPatient[2]\n",
    "            except:\n",
    "                ct_z, cta_p_z, cta_p1_z = ct[i].SliceLocation, cta[p].SliceLocation, cta[p+1].SliceLocation\n",
    "            if (cta_p_z-ct_z) * (cta_p1_z-ct_z) <= 0:\n",
    "                p = p if abs(cta_p_z-ct_z) < abs(cta_p1_z-ct_z) else p+1\n",
    "                with open(os.path.join(cta_root_path, name, 'labels', f'{name}_{p:04d}.txt'), 'r') as sf:\n",
    "                    cta_labels = sf.readlines()\n",
    "                cta_img = cta[p].pixel_array.astype(np.int16)\n",
    "                cta_intercept = cta[p].RescaleIntercept\n",
    "                cta_slope = cta[p].RescaleSlope\n",
    "                if cta_slope != 1:\n",
    "                    cta_img = (cta_slope * cta_img.astype(np.float64)).astype(np.int16)\n",
    "                cta_img += np.int16(cta_intercept)\n",
    "                cta_img = np.clip(cta_img, -100, 500)\n",
    "                cta_img = ((cta_img+100)/600*255).astype(np.uint8)\n",
    "                break\n",
    "            p += 1\n",
    "        else:\n",
    "            raise Exception(f'CT图像范围超出了CTA图像范围：{ct_z},{cta_p_z},{cta_p1_z}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "a9eabd9b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Errno 2] No such file or directory: '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/pred_labels/hujiaguo-Im34-58_0019.txt'\n",
      "[Errno 2] No such file or directory: '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/pred_labels/hujiaguo-Im34-58_0020.txt'\n",
      "[Errno 2] No such file or directory: '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/pred_labels/shengfumei-Im31-132_0020.txt'\n",
      "[Errno 2] No such file or directory: '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/pred_labels/shengfumei-Im31-132_0122.txt'\n",
      "[Errno 2] No such file or directory: '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/pred_labels/shengfumei-Im31-132_0123.txt'\n",
      "[Errno 2] No such file or directory: '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/pred_labels/shengfumei-Im31-132_0125.txt'\n",
      "[Errno 2] No such file or directory: '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/pred_labels/shengfumei-Im31-132_0126.txt'\n",
      "[Errno 2] No such file or directory: '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/pred_labels/shengfumei-Im31-132_0127.txt'\n",
      "[Errno 2] No such file or directory: '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/pred_labels/shengfumei-Im31-132_0128.txt'\n",
      "[Errno 2] No such file or directory: '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/pred_labels/shengfumei-Im31-132_0129.txt'\n",
      "[Errno 2] No such file or directory: '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/pred_labels/shengfumei-Im31-132_0130.txt'\n",
      "[Errno 2] No such file or directory: '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/pred_labels/shengfumei-Im31-132_0131.txt'\n",
      "[Errno 2] No such file or directory: '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/pred_labels/shengfumei-Im31-132_0132.txt'\n",
      "[Errno 2] No such file or directory: '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/pred_labels/shengfumei-Im31-132_0133.txt'\n",
      "[Errno 2] No such file or directory: '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/pred_labels/shengfumei-Im31-132_0134.txt'\n",
      "[Errno 2] No such file or directory: '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/pred_labels/shengfumei-Im31-132_0135.txt'\n",
      "[Errno 2] No such file or directory: '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/pred_labels/xujinbiao-Im37-146_0146.txt'\n",
      "[Errno 2] No such file or directory: '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/pred_labels/zhengsiqing-Im35-106_0028.txt'\n",
      "[Errno 2] No such file or directory: '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/pred_labels/zhengsiqing-Im35-106_0029.txt'\n",
      "[Errno 2] No such file or directory: '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/pred_labels/zhengsiqing-Im35-106_0030.txt'\n",
      "[Errno 2] No such file or directory: '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/pred_labels/zhengsiqing-Im35-106_0159.txt'\n",
      "[Errno 2] No such file or directory: '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/pred_labels/zhengsiqing-Im35-106_0160.txt'\n",
      "branch recall:0.9914529914529915, other recall:0.9445544554455445\n"
     ]
    }
   ],
   "source": [
    "val_label_path = '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/labels/val/'\n",
    "result_label_path = '/nfs3-p1/zsxm/dataset/aorta_yolo_ct/pred_labels/'\n",
    "\n",
    "branch_count, branch_corr, other_count, other_corr = 0,0,0,0\n",
    "\n",
    "def iou(pred, target):\n",
    "    sp, st = pred[2]*pred[3], target[2]*target[3]\n",
    "    px1, py1, px2, py2 = pred[0]-pred[2]/2, pred[1]-pred[3]/2, pred[0]+pred[2]/2, pred[1]+pred[3]/2\n",
    "    tx1, ty1, tx2, ty2 = target[0]-target[2]/2, target[1]-target[3]/2, target[0]+target[2]/2, target[1]+target[3]/2\n",
    "    ix1, iy1, ix2, iy2 = max(px1, tx1), max(py1, ty1), min(px2, tx2), min(py2, ty2)\n",
    "    if ix1>ix2 or iy1>iy2:\n",
    "        return 0.0\n",
    "    si = max((ix2-ix1)*(iy2-iy1), 0)\n",
    "    i_o_u = si/(sp+st-si)\n",
    "    assert 0.0<=i_o_u<=1.0\n",
    "    return i_o_u\n",
    "\n",
    "for label in os.listdir(val_label_path):\n",
    "    with open(os.path.join(val_label_path, label), 'r') as vf:\n",
    "        val_lines = vf.readlines()\n",
    "    assert 0<=len(val_lines)<=2\n",
    "    if len(val_lines) == 0:\n",
    "        continue\n",
    "        \n",
    "    if len(val_lines) == 1:\n",
    "        other_count += 1\n",
    "        tx, ty, tw, th = list(map(lambda x: float(x), val_lines[0].split()))[1:]\n",
    "        try:\n",
    "            with open(os.path.join(result_label_path, label), 'r') as rf:\n",
    "                result_lines = rf.readlines()\n",
    "        except Exception as e:\n",
    "            print(e)\n",
    "            continue\n",
    "        for i in range(len(result_lines)):\n",
    "            px, py, pw, ph = list(map(lambda x: float(x), result_lines[i].split()))[1:]\n",
    "            if iou((px, py, pw, ph), (tx, ty, tw, th)) > 0.5:\n",
    "                other_corr += 1\n",
    "                break\n",
    "    else:\n",
    "        other_count += 1\n",
    "        branch_count += 1\n",
    "        line0, line1 = val_lines[0].split(), val_lines[1].split()\n",
    "        line0_f, line1_f = list(map(lambda x: float(x), line0)), list(map(lambda x: float(x), line1))\n",
    "        if line0_f[2] > line1_f[2]:\n",
    "            aorta, branch = line0_f, line1_f\n",
    "        else:\n",
    "            aorta, branch = line1_f, line0_f\n",
    "        atx, aty, atw, ath = aorta[1:]\n",
    "        btx, bty, btw, bth = branch[1:]\n",
    "        try:\n",
    "            with open(os.path.join(result_label_path, label), 'r') as rf:\n",
    "                result_lines = rf.readlines()\n",
    "        except:\n",
    "            continue\n",
    "        for i in range(len(result_lines)):\n",
    "            px, py, pw, ph = list(map(lambda x: float(x), result_lines[i].split()))[1:]\n",
    "            if iou((px, py, pw, ph), (atx, aty, atw, ath)) > 0.5:\n",
    "                other_corr += 1\n",
    "                break\n",
    "        for i in range(len(result_lines)):\n",
    "            px, py, pw, ph = list(map(lambda x: float(x), result_lines[i].split()))[1:]\n",
    "            if iou((px, py, pw, ph), (btx, bty, btw, bth)) > 0.5:\n",
    "                branch_corr += 1\n",
    "                break\n",
    "                \n",
    "print(f'branch recall:{branch_corr/branch_count}, other recall:{other_corr/other_count}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ff6ce89",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
